#!/usr/bin/python2
# coding: utf8
import rdkit.Chem as Chem
from rdkit.Chem import Draw, AllChem
import os, sys, cStringIO
from fm_parser import exact_mass

class Mol2pic():
    def __init__(self, mol_string = '', width = 600, height = 600):
        self.size = (width, height)
        self.mol = None
        self.readers = {
                'smi': Chem.MolFromSmiles,
                'smiles': Chem.MolFromSmiles,
                'sma': Chem.MolFromSmarts,
                'smarts': Chem.MolFromSmarts,
                'mol': Chem.MolFromMolBlock,
                'mol2': Chem.MolFromMol2Block,
                'inchi': Chem.MolFromInchi,
                }
        if mol_string:
            self.set_mol(mol_string)

    def set_mol(self, string, fmt = 'mol'):
        #string = open('d.mol', 'r').write()
        try:
            string = string.encode('utf8')
        except: pass
        if os.path.exists(string):
            string = open(string, 'r').read()
        self.mol = self.readers[fmt](string)
        self.mol = Chem.RemoveHs(self.mol, True)

    def to_pic(self, out_fmt = 'svg', auto_layout = False, filename = None):
        if auto_layout:
            AllChem.Compute2DCoords(self.mol)
        if filename:
            try:
                Draw.MolToFile(self.mol, filename, imageType = out_fmt)
                return filename
            except Exception, e:
                print e
        filename = cStringIO.StringIO()
        Draw.MolToFile(self.mol, filename, imageType = out_fmt)
        svg = filename.getvalue()
        filename.close()
        return svg

    def to_mol(self, auto_layout = False):
        if auto_layout:
            AllChem.Compute2DCoords(self.mol)
        return Chem.MolToMolBlock(self.mol)

    def to_smi(self):
        return Chem.MolToSmiles(self.mol)

def sub_filter(sub, mol):
    ptn = Chem.MolFromSmiles(sub)
    if ptn and mol.HasSubstructMatch(ptn):
        return True
    return False

def filter_sdf(sdf, filters = []):
    #open('a.sdf', 'w').write(sdf)
    results = []
    if not filters:
        num_all = sdf.count('M  END')
        if num_all == 0:
            num_all = sdf.count('$$$$')
        return {'sdf': sdf, 'num_all': num_all, 'num_hits': num_all}
    if os.path.exists(sdf):
        sdffile = sdf
    else:
        # store sdf to temp file
        sdffile = tempfile.mkstemp(text = True)[1]
        open(sdffile, 'w').write(sdf)
    mols = Chem.SDMolSupplier(sdffile)
    hits = []
    for num_all, mol in enumerate(mols):
        for filter in filters:
            if filter(mol):
                hits.append(mol)
    num_all += 1
    sio = cStringIO.StringIO()
    w = Chem.SDWriter(sio)
    for mol in hits:
        w.write(mol)
    w.flush()
    w.close()
    # remove the temp sdffile
    try:
        os.remove(sdffile)
    except:
        pass
    return {'sdf': sio.getvalue(), 'num_all': num_all, 'num_hits': len(hits)}

if __name__ == '__main__':
    smi = 'CC(=O)NCCC1=CNc2c1cc(OC)cc2'
    mp = Mol2pic()
    mp.set_mol(smi, 'smi')
    print mp.to_pic(filename = 'd.svg')

